'''
Author:
Date: 2022-01-06 18:34:00
LastEditTime: 2022-12-08 15:12:54
LastEditors: Please set LastEditors
'''

import sys
import time
import cv2

sys.path.append("D:\AI\YOLOX") #YOLOX安装目录

from loguru import logger
from pill_handler import YoloxHandler


file_idx = 0
save_interval = 1
time_temp = time.time()


if __name__ == "__main__":
    # 模型
    exp_file = "d:/AI/YOLOX/exps/example/custom/yolox_s.py"

    # 参数
    ckpt = "D:/AI/YOLOX/YOLOX_outputs/yolox_s/last_epoch_ckpt.pth"

    handle = YoloxHandler(exp_file, ckpt, False)

    camera_id = 0
    cap = cv2.VideoCapture(camera_id, cv2.CAP_DSHOW)

    if not cap.isOpened():
        logger.error("can not open camera {}".format(camera_id))

    while cap.isOpened():
        success, frame = cap.read()
        assert success, 'no frame read'

        outputs, img_info = handle.predict(frame)  # 预测
         
        frame=handle.visual(outputs[0],img_info,cls_conf=0.5) #画框

        cv2.imshow('my_window', frame)
        ch = cv2.waitKey(1)
        if ch == 27 or ch == ord("q") or ch == ord("Q"):   #按esc或者q退出
            break
    cap.release()

    cv2.destroyAllWindows()
